package ssh

import (
	"github.com/kataras/golog"
	"golang.org/x/crypto/ssh"
	"net"
	"sdn-executor/config"
	"sdn-executor/utils/erros"
	"strconv"
	"time"
)

type Client struct {
	*Config

	tcpConn   *net.TCPConn
	sshClient *ssh.Client
}

func NewClient(config *Config) (*Client, *erros.Error) {

	client, conn, err := DialSSH(config)
	if err != nil {
		golog.Error("DialSSH failed")
		return nil, err
	}

	return &Client{
		Config:    config,
		tcpConn:   conn,
		sshClient: client,
	}, nil
}

func NewSimpleClient(device string, port int16, user string, pwd string) (*Client, *erros.Error) {
	sshConfig := &Config{
		Device:   device,
		Port:     port,
		User:     user,
		Password: pwd,

		ConnTimeout:  config.C().SSH.ConnTimeoutMs * time.Millisecond,
		ReadTimeout:  config.C().SSH.ReadTimeoutMs * time.Millisecond,
		WriteTimeout: config.C().SSH.WriteTimeoutMs * time.Millisecond,
	}

	return NewClient(sshConfig)
}

func DialSSH(config *Config) (*ssh.Client, *net.TCPConn, *erros.Error) {

	// ssh配置
	sshConfig := &ssh.ClientConfig{
		User:            config.User,
		Auth:            []ssh.AuthMethod{ssh.Password(config.Password)},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         config.ConnTimeout,
	}

	// 建立tcp连接
	addr := config.Device + ":" + strconv.Itoa(int(config.Port))
	conn, err := net.DialTimeout("tcp", addr, sshConfig.Timeout)
	if err != nil {
		golog.Error("net.DialTimeout failed.")
		return nil, nil, erros.New(erros.CodeInternalError, "DialTimeout")
	}

	// 建立ssh连接
	deadlineConn := &deadlineConn{TCPConn: conn.(*net.TCPConn), readTimeout: config.ReadTimeout, writeTimeout: config.WriteTimeout}
	c, chans, reqs, err := ssh.NewClientConn(deadlineConn, addr, sshConfig)
	if err != nil {
		golog.Error("ssh.NewClientConn failed.")
		return nil, nil, erros.New(erros.CodeInternalError, "NewClientConn")
	}

	// 创建ssh Client
	return ssh.NewClient(c, chans, reqs), deadlineConn.TCPConn, nil
}

func (c *Client) NewSession(subsystem Subsystem) (*Session, *erros.Error) {

	// 建立新会话
	sshSession, err := c.sshClient.NewSession()
	if err != nil {
		golog.Errorf("new sshSession failed. device: %s, err: %+v", c.Device, err)
		return nil, erros.New(erros.CodeInternalError, "create sshSession failed")
	}

	// 设置tty
	if subsystem == SubsystemShell {
		if e := setSSHTerminal(sshSession); e != nil {
			golog.Error("request pty failed. device: %s, err: %+v", c.Device, e)
			sshSession.Close()
			return nil, e
		}
	}

	// 打开shell 通道
	reader, err := sshSession.StdoutPipe()
	if err != nil {
		sshSession.Close()
		return nil, erros.New(erros.CodeInternalError, "open StdoutPipe failed")
	}
	writer, err := sshSession.StdinPipe()
	if err != nil {
		sshSession.Close()
		return nil, erros.New(erros.CodeInternalError, "open StdinPipe failed")
	}
	if subsystem == SubsystemShell {
		err = sshSession.Shell()
		if err != nil {
			golog.Infof("open shell failed. device: %s, error: %s", c.Device, err.Error())
			sshSession.Close()
			return nil, erros.New(erros.CodeInternalError, "open Shell failed")
		}
	} else if subsystem == SubsystemNetconf {
		err = sshSession.RequestSubsystem(string(subsystem))
		if err != nil {
			golog.Errorf("open netconf subsystem failed. device: %s, err: %+v", c.Device, err)
			sshSession.Close()
			return nil, erros.New(erros.CodeInternalError, "open netconf subsystem failed")
		}
	} else {
		sshSession.Close()
		return nil, erros.New(erros.CodeUnsupported, "unsupported")
	}

	return &Session{
		Client:     c,
		subsystem:  subsystem,
		sshSession: sshSession,
		reader:     reader,
		writer:     writer,
	}, nil
}

func (c *Client) Close() {
	if c != nil && c.sshClient != nil {
		c.sshClient.Close()
	}
}

// deadlineConn 带读写超时的connection
type deadlineConn struct {
	*net.TCPConn
	readTimeout  time.Duration
	writeTimeout time.Duration
}

func (c *deadlineConn) Read(b []byte) (n int, err error) {
	if c.readTimeout > 0 {
		err = c.SetReadDeadline(time.Now().Add(c.readTimeout))
	}
	return c.TCPConn.Read(b)
}

func (c *deadlineConn) Write(b []byte) (n int, err error) {
	if c.writeTimeout > 0 {
		err = c.SetWriteDeadline(time.Now().Add(c.writeTimeout))
	}
	return c.TCPConn.Write(b)
}

func setSSHTerminal(session *ssh.Session) *erros.Error {
	modes := ssh.TerminalModes{
		ssh.ECHO:          0,      // 禁用回显（0禁用，1启动）
		ssh.TTY_OP_ISPEED: 115200, // input speed = 115.2kbaud
		ssh.TTY_OP_OSPEED: 115200, // output speed = 115.2kbaud
	}
	err := session.RequestPty("vt102", 512, 1024, modes)
	if err != nil {
		return erros.New(erros.CodeInternalError, "RequestPty failed.")
	}
	return nil
}
